import ast
import json
import javalang
import esprima
import argparse
import os
import pickle
import time
from pathlib import Path
import subprocess
import torch.multiprocessing as mp
import os.path as osp
from datasets import load_dataset, load_from_disk
import sys
import pandas as pd  # 新增导入

sys.path.insert(0, "/data_ext/ref_code/LocAgent")
from dependency_graph.build_graph import build_graph, VERSION
from util.benchmark.setup_repo import setup_repo

VERSION = "v2.3"
NODE_TYPE_DIRECTORY = "directory"
NODE_TYPE_FILE = "file"
NODE_TYPE_CLASS = "class"
NODE_TYPE_FUNCTION = "function"

EDGE_TYPE_CONTAINS = "contains"
EDGE_TYPE_INHERITS = "inherits"
EDGE_TYPE_INVOKES = "invokes"
EDGE_TYPE_IMPORTS = "imports"

VALID_NODE_TYPES = [
    NODE_TYPE_DIRECTORY,
    NODE_TYPE_FILE,
    NODE_TYPE_CLASS,
    NODE_TYPE_FUNCTION,
]
VALID_EDGE_TYPES = [
    EDGE_TYPE_CONTAINS,
    EDGE_TYPE_INHERITS,
    EDGE_TYPE_INVOKES,
    EDGE_TYPE_IMPORTS,
]


class JavaAnalyzer:
    def __init__(self, filename, source_code):
        self.filename = filename
        self.source_code = source_code
        self.source_lines = source_code.split("\n")
        self.nodes = []
        self.edges = []
        self.package_name = ""
        self.imports = []
        self.current_class = None

    def analyze(self, tree):
        """分析Java语法树并提取节点和边"""
        if hasattr(tree, "package") and tree.package:
            self.package_name = tree.package.name

        # 提取import信息
        if hasattr(tree, "imports"):
            for import_decl in tree.imports:
                self.imports.append(import_decl.path)

        # 分析类型声明
        if hasattr(tree, "types"):
            for type_decl in tree.types:
                self._visit_type(type_decl)

        return {"nodes": self.nodes, "edges": self.edges}

    def _get_line_number(self, node):
        """安全地获取节点的行号"""
        try:
            if hasattr(node, "position") and node.position:
                return node.position.line
        except:
            pass
        return 0

    def _get_end_line_number(self, node):
        """获取节点的结束行号"""
        try:
            max_line = self._get_line_number(node)
            if hasattr(node, "body") and node.body:
                for child in node.body:
                    child_line = self._get_line_number(child)
                    if child_line > max_line:
                        max_line = child_line
            return max_line if max_line > 0 else None
        except:
            return None

    def _extract_code_block(self, start_line, end_line):
        """提取指定行范围的代码"""
        if start_line <= 0 or start_line > len(self.source_lines):
            return ""

        if end_line is None or end_line <= 0:
            end_line = start_line

        start_idx = start_line - 1
        end_idx = min(end_line, len(self.source_lines))

        return "\n".join(self.source_lines[start_idx:end_idx])

    def _visit_type(self, node):
        """访问类型声明节点"""
        if isinstance(node, javalang.tree.ClassDeclaration):
            self._visit_class(node)
        elif isinstance(node, javalang.tree.InterfaceDeclaration):
            self._visit_interface(node)
        elif isinstance(node, javalang.tree.EnumDeclaration):
            self._visit_enum(node)

    def _visit_class(self, node):
        """访问类声明"""
        class_name = node.name
        if self.package_name:
            full_class_name = f"{self.package_name}.{class_name}"
        else:
            full_class_name = class_name

        start_line = self._get_line_number(node)
        end_line = self._get_end_line_number(node)
        code_block = self._extract_code_block(start_line, end_line)

        class_node = {
            "name": full_class_name,
            "parent_type": None,
            "type": NODE_TYPE_CLASS,
            "code": code_block if code_block else f"class {class_name}",
            "start_line": start_line,
            "end_line": end_line,
            "file_path": self.filename,
            "modifiers": [mod for mod in node.modifiers] if node.modifiers else [],
        }

        self.nodes.append(class_node)
        self.current_class = full_class_name

        # 处理继承关系
        if hasattr(node, "extends") and node.extends:
            parent_class = node.extends.name
            self._add_inheritance_edge(full_class_name, parent_class)

        # 处理实现的接口
        if hasattr(node, "implements") and node.implements:
            for interface in node.implements:
                interface_name = interface.name
                self._add_inheritance_edge(full_class_name, interface_name)

        # 处理类成员
        if hasattr(node, "body"):
            for member in node.body:
                if isinstance(member, javalang.tree.MethodDeclaration):
                    self._visit_method(member, class_name)
                elif isinstance(member, javalang.tree.ConstructorDeclaration):
                    self._visit_constructor(member, class_name)
                elif isinstance(member, javalang.tree.FieldDeclaration):
                    self._visit_field(member, class_name)

    def _visit_interface(self, node):
        """访问接口声明"""
        interface_name = node.name
        if self.package_name:
            full_interface_name = f"{self.package_name}.{interface_name}"
        else:
            full_interface_name = interface_name

        start_line = self._get_line_number(node)
        end_line = self._get_end_line_number(node)
        code_block = self._extract_code_block(start_line, end_line)

        interface_node = {
            "name": full_interface_name,
            "parent_type": None,
            "type": NODE_TYPE_CLASS,
            "code": code_block if code_block else f"interface {interface_name}",
            "start_line": start_line,
            "end_line": end_line,
            "file_path": self.filename,
            "is_interface": True,
            "modifiers": [mod for mod in node.modifiers] if node.modifiers else [],
        }

        self.nodes.append(interface_node)
        self.current_class = full_interface_name

        # 处理接口继承
        if hasattr(node, "extends") and node.extends:
            for parent_interface in node.extends:
                parent_name = parent_interface.name
                self._add_inheritance_edge(full_interface_name, parent_name)

        # 处理接口方法
        if hasattr(node, "body"):
            for member in node.body:
                if isinstance(member, javalang.tree.MethodDeclaration):
                    self._visit_method(member, interface_name)

    def _visit_enum(self, node):
        """访问枚举声明"""
        enum_name = node.name
        if self.package_name:
            full_enum_name = f"{self.package_name}.{enum_name}"
        else:
            full_enum_name = enum_name

        start_line = self._get_line_number(node)
        end_line = self._get_end_line_number(node)
        code_block = self._extract_code_block(start_line, end_line)
        enum_node = {
            "name": full_enum_name,
            "parent_type": None,
            "type": NODE_TYPE_CLASS,
            "code": code_block if code_block else f"enum {enum_name}",
            "start_line": start_line,
            "end_line": end_line,
            "file_path": self.filename,
            "is_enum": True,
            "modifiers": [mod for mod in node.modifiers] if node.modifiers else [],
        }

        self.nodes.append(enum_node)
        self.current_class = full_enum_name

    def _visit_method(self, node, parent_class):
        """访问方法声明"""
        method_name = node.name
        if self.package_name:
            full_method_name = f"{self.package_name}.{parent_class}.{method_name}"
        else:
            full_method_name = f"{parent_class}.{method_name}"

        start_line = self._get_line_number(node)
        end_line = self._get_end_line_number(node)
        code_block = self._extract_code_block(start_line, end_line)

        # 获取参数信息
        parameters = []
        if hasattr(node, "parameters") and node.parameters:
            for param in node.parameters:
                param_info = {
                    "name": param.name,
                    "type": str(param.type) if param.type else "unknown",
                }
                parameters.append(param_info)

        # 获取返回类型
        return_type = str(node.return_type) if node.return_type else "void"

        method_node = {
            "name": full_method_name,
            "parent_type": NODE_TYPE_CLASS,
            "type": NODE_TYPE_FUNCTION,
            "code": code_block if code_block else f"method {method_name}",
            "start_line": start_line,
            "end_line": end_line,
            "file_path": self.filename,
            "parameters": parameters,
            "return_type": return_type,
            "modifiers": [mod for mod in node.modifiers] if node.modifiers else [],
        }

        self.nodes.append(method_node)

        # 添加包含关系边
        parent_full_name = (
            f"{self.package_name}.{parent_class}" if self.package_name else parent_class
        )
        self._add_contains_edge(parent_full_name, full_method_name)

        # 分析方法体中的方法调用
        if hasattr(node, "body") and node.body:
            self._analyze_method_calls(node.body, full_method_name)

    def _visit_constructor(self, node, parent_class):
        """访问构造函数声明"""
        constructor_name = node.name
        if self.package_name:
            full_constructor_name = (
                f"{self.package_name}.{parent_class}.{constructor_name}"
            )
        else:
            full_constructor_name = f"{parent_class}.{constructor_name}"

        start_line = self._get_line_number(node)
        end_line = self._get_end_line_number(node)
        code_block = self._extract_code_block(start_line, end_line)

        # 获取参数信息
        parameters = []
        if hasattr(node, "parameters") and node.parameters:
            for param in node.parameters:
                param_info = {
                    "name": param.name,
                    "type": str(param.type) if param.type else "unknown",
                }
                parameters.append(param_info)

        constructor_node = {
            "name": full_constructor_name,
            "parent_type": NODE_TYPE_CLASS,
            "type": NODE_TYPE_FUNCTION,
            "code": code_block if code_block else f"constructor {constructor_name}",
            "start_line": start_line,
            "end_line": end_line,
            "file_path": self.filename,
            "parameters": parameters,
            "is_constructor": True,
            "modifiers": [mod for mod in node.modifiers] if node.modifiers else [],
        }

        self.nodes.append(constructor_node)

        # 添加包含关系边
        parent_full_name = (
            f"{self.package_name}.{parent_class}" if self.package_name else parent_class
        )
        self._add_contains_edge(parent_full_name, full_constructor_name)

    def _visit_field(self, node, parent_class):
        """访问字段声明"""
        if hasattr(node, "declarators"):
            for declarator in node.declarators:
                field_name = declarator.name
                if self.package_name:
                    full_field_name = f"{self.package_name}.{parent_class}.{field_name}"
                else:
                    full_field_name = f"{parent_class}.{field_name}"

                start_line = self._get_line_number(node)
                field_type = str(node.type) if node.type else "unknown"

                field_node = {
                    "name": full_field_name,
                    "parent_type": NODE_TYPE_CLASS,
                    "type": "field",
                    "field_type": field_type,
                    "start_line": start_line,
                    "file_path": self.filename,
                    "modifiers": (
                        [mod for mod in node.modifiers] if node.modifiers else []
                    ),
                }

                self.nodes.append(field_node)

    def _analyze_method_calls(self, body, caller_method):
        """分析方法体中的方法调用"""
        try:
            for node in body:
                if hasattr(node, "children"):
                    for child in node.children:
                        if isinstance(child, javalang.tree.MethodInvocation):
                            called_method = child.member
                            self._add_invocation_edge(caller_method, called_method)
        except Exception as e:
            # 方法调用分析可能比较复杂，这里简化处理
            pass

    def _add_inheritance_edge(self, child_class, parent_class):
        """添加继承关系边"""
        edge = {
            "source": child_class,
            "target": parent_class,
            "type": EDGE_TYPE_INHERITS,
        }
        self.edges.append(edge)

    def _add_contains_edge(self, container, contained):
        """添加包含关系边"""
        edge = {"source": container, "target": contained, "type": EDGE_TYPE_CONTAINS}
        self.edges.append(edge)

    def _add_invocation_edge(self, caller, callee):
        """添加方法调用边"""
        edge = {"source": caller, "target": callee, "type": EDGE_TYPE_INVOKES}
        self.edges.append(edge)

    def _add_import_edge(self, importer, imported):
        """添加导入关系边"""
        edge = {"source": importer, "target": imported, "type": EDGE_TYPE_IMPORTS}
        self.edges.append(edge)


def analyze_file_java(filepath):
    """分析Java文件"""
    try:
        with open(filepath, "r", encoding="utf-8") as file:
            code = file.read()

        if filepath.endswith(".java"):
            tree = javalang.parse.parse(code)
            analyzer = JavaAnalyzer(filepath, code)
            result = analyzer.analyze(tree)
            return result["nodes"], result["edges"]
        elif filepath.endswith(".js"):
            # JavaScript分析器可以在这里实现
            tree = esprima.parseScript(code)
            return [], []
        else:
            return [], []
    except Exception as e:
        print(f"Error parsing {filepath}: {e}")
        return [], []


def analyze_java_project(project_path):
    """分析整个Java项目"""
    all_nodes = []
    all_edges = []

    # 遍历项目目录
    for root, dirs, files in os.walk(project_path):
        for file in files:
            if file.endswith(".java"):
                file_path = os.path.join(root, file)
                try:
                    nodes, edges = analyze_file_java(file_path)
                    all_nodes.extend(nodes)
                    all_edges.extend(edges)
                except Exception as e:
                    print(f"Error analyzing {file_path}: {e}")

    return {"nodes": all_nodes, "edges": all_edges}


def edges_to_dataframe(edges):
    """将边信息转换为DataFrame"""
    if not edges:
        return pd.DataFrame()

    # 处理复杂字段，将列表和字典转换为字符串
    processed_edges = []
    for edge in edges:
        processed_edge = {}
        for key, value in edge.items():
            if isinstance(value, (list, dict)):
                processed_edge[key] = json.dumps(value, ensure_ascii=False)
            else:
                processed_edge[key] = value
        processed_edges.append(processed_edge)

    df = pd.DataFrame(processed_edges)
    return df


def save_edges_to_excel(edges, output_path):
    """将边信息保存为Excel文件"""
    df = edges_to_dataframe(edges)
    if not df.empty:
        df.to_excel(output_path, index=False, engine="openpyxl")
        print(f"Saved {len(df)} edges to {output_path}")
    else:
        print("No edges to save")


def nodes_to_dataframe(nodes):
    """将节点信息转换为DataFrame"""
    if not nodes:
        return pd.DataFrame()

    # 处理复杂字段，将列表和字典转换为字符串
    processed_nodes = []
    for node in nodes:
        processed_node = {}
        for key, value in node.items():
            if isinstance(value, (list, dict)):
                processed_node[key] = json.dumps(value, ensure_ascii=False)
            else:
                processed_node[key] = value
        processed_nodes.append(processed_node)

    df = pd.DataFrame(processed_nodes)
    return df


def save_nodes_to_csv(nodes, output_path):
    """将节点信息保存为CSV文件"""
    df = nodes_to_dataframe(nodes)
    if not df.empty:
        df.to_csv(output_path, index=False, encoding="utf-8")
        print(f"Saved {len(df)} nodes to {output_path}")
    else:
        print("No nodes to save")


def save_nodes_to_excel(nodes, output_path):
    """将节点信息保存为Excel文件"""
    df = nodes_to_dataframe(nodes)
    if not df.empty:
        df.to_excel(output_path, index=False, engine="openpyxl")
        print(f"Saved {len(df)} nodes to {output_path}")
    else:
        print("No nodes to save")


def list_folders(path):
    return [p.name for p in Path(path).iterdir() if p.is_dir()]


def run(rank, repo_name, repo_path, out_path, download_repo=False, instance_data=None):
    """处理单个仓库"""
    while True:
        output_file = f"{osp.join(out_path, repo_name)}.pkl"
        repo_dir = osp.join(repo_path, repo_name)

        print(f"[{rank}] Start process {repo_name}")
        try:
            # 使用改进的Java分析器
            if os.path.exists(repo_dir):
                result = analyze_java_project(repo_dir)

                # 保存原始pickle文件
                with open(output_file, "wb") as f:
                    pickle.dump(result, f)

                # 新增：保存节点信息到DataFrame和CSV文件
                if result["nodes"]:
                    # 保存为CSV
                    # csv_output_file = f"{osp.join(out_path, repo_name)}_nodes.csv"
                    # save_nodes_to_csv(result["nodes"], csv_output_file)

                    # 可选：保存为Excel
                    excel_output_file = f"{osp.join(out_path, repo_name)}_nodes.xlsx"
                    save_nodes_to_excel(result["nodes"], excel_output_file)

                    # 打印DataFrame信息
                    df = nodes_to_dataframe(result["nodes"])
                    print(f"[{rank}] DataFrame shape: {df.shape}")
                    print(f"[{rank}] DataFrame columns: {list(df.columns)}")
                if result["edges"]:
                    edges_excel_output_file = (
                        f"{osp.join(out_path, repo_name)}_edges.xlsx"
                    )
                    save_edges_to_excel(result["edges"], edges_excel_output_file)

                    # 打印边的DataFrame信息
                    edges_df = edges_to_dataframe(result["edges"])
                    print(f"[{rank}] Edges DataFrame shape: {edges_df.shape}")
                    print(f"[{rank}] Edges DataFrame columns: {list(edges_df.columns)}")

                print(
                    f"[{rank}] Processed {repo_name} - Found {len(result['nodes'])} nodes and {len(result['edges'])} edges"
                )
            else:
                print(f"[{rank}] Repository directory not found: {repo_dir}")
        except Exception as e:
            print(f"[{rank}] Error processing {repo_name}: {e}")
        break


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=str, default="pp_data")
    parser.add_argument("--split", type=str, default="test")
    parser.add_argument("--num_processes", type=int, default=1)
    parser.add_argument(
        "--download_repo",
        action="store_true",
        default=False,
        help="Whether to download the codebase to `repo_path` before indexing.",
    )
    parser.add_argument(
        "--repo_path",
        type=str,
        default="/data_ext/ref_code/LocAgent/data/data/pp_data",
        help="The directory where you plan to pull or have already pulled the codebase.",
    )
    parser.add_argument(
        "--index_dir",
        type=str,
        default="index_data",
        help="The base directory where the generated graph index will be saved.",
    )
    parser.add_argument(
        "--instance_id_path",
        type=str,
        default="",
        help="Path to a file containing a list of selected instance IDs.",
    )
    args = parser.parse_args()

    dataset_name = args.dataset.split("/")[-1]
    args.index_dir = f"{args.index_dir}/{dataset_name}/graph_index_{VERSION}_std/"
    os.makedirs(args.index_dir, exist_ok=True)

    # 加载选定的仓库实例ID和实例数据
    if args.download_repo:
        selected_instance_data = {}
        bench_data = load_from_disk(args.dataset)["test"]
        if args.instance_id_path and osp.exists(args.instance_id_path):
            with open(args.instance_id_path, "r") as f:
                repo_folders = json.loads(f.read())
            for instance in bench_data:
                if instance["instance_id"] in repo_folders:
                    selected_instance_data[instance["instance_id"]] = instance
        else:
            repo_folders = []
            for instance in bench_data:
                repo_folders.append(instance["instance_id"])
                selected_instance_data[instance["instance_id"]] = instance
    else:
        if args.instance_id_path and osp.exists(args.instance_id_path):
            with open(args.instance_id_path, "r") as f:
                repo_folders = json.loads(f.read())
        else:
            repo_folders = list_folders(args.repo_path)
        selected_instance_data = None

    os.makedirs(args.repo_path, exist_ok=True)

    # 创建共享队列并添加仓库
    manager = mp.Manager()
    queue = manager.Queue()
    for repo in repo_folders:
        queue.put(repo)

    start_time = time.time()

    repo_name = "/data_ext/ref_code/LocAgent/data/data/pp_data/pingp-auto-test"
    run(
        0,
        repo_name,
        args.repo_path,
        args.index_dir,
        args.download_repo,
        selected_instance_data,
    )

    end_time = time.time()
    print(f"Total Execution time = {end_time - start_time:.3f}s")

    # 测试单个文件
    # if len(repo_folders) == 1:
    #     test_file = "/data_ext/ref_code/LocAgent/data/data/pp_data/pingp-auto-test/auto-test-b2b/src/test/java/auto/b2b/check/CheckBusinessAccount.java"
    #     if os.path.exists(test_file):
    #         nodes, edges = analyze_file_java(test_file)
    #         print(
    #             f"\nTest file analysis - Found {len(nodes)} nodes and {len(edges)} edges:"
    #         )
    #         for node in nodes[:3]:  # 只显示前3个节点
    #             print(json.dumps(node, ensure_ascii=False, indent=2))
    #
    #         # 新增：测试DataFrame功能
    #         if nodes:
    #             df = nodes_to_dataframe(nodes)
    #             print(f"\nDataFrame shape: {df.shape}")
    #             print(f"DataFrame columns: {list(df.columns)}")
    #             print("\nFirst few rows:")
    #             print(df.head())
